import os

from typing import Optional, List
from fastapi import FastAPI, Form

import base64
import cv2
import requests
from face_embedding.register_pkl import register_pkl_embedding, re_register_pkl_embedding
from face_embedding.one2one_test import one2onne_embedding
from face_embedding.face_crop import FaceRetinaModel
from face_embedding.utils import argsutils
from face_embedding.face_feature_thread import FaceFeatureThread
from utils import dictProperties, dbpool_util

import numpy as np

import traceback
from datetime import datetime
import json
import re
import time

app = FastAPI()

faiss_server_url = dictProperties.get('faiss_server')

args = argsutils.Facemodel(dictProperties).args
retina_model = FaceRetinaModel(dictProperties)
face_feature = FaceFeatureThread(args, dictProperties)

def get_db():
    return dbpool_util

def query_db(query, args=(), one=False):
    # print('sql_s1', query, file=sys.stdout)
    result = get_db().execute_query(query, args)
    return result

def execute_db(sql, args=()):
    return get_db().execute_iud(sql, args)

def get_user_info(user_code):
    sql_s = "select id, user_code, group_id,username,user_remark,embedding from t_user_info where user_code='{user_code}';".format(
        user_code=user_code)
    user_info = query_db(sql_s)
    if user_info and len(user_info) > 0:
        res =[]
        for emb_lst in json.loads(user_info[0][5]):
            res.append(np.array(emb_lst, dtype=np.float32))

        return {
            'userCode': user_info[0][1],
            'groupId': user_info[0][2],
            'username': user_info[0][3],
            'embedding': res
        }
    else:
        return None

def insert_user_info(user_code, group_id, user_name, embedding):
    embedding_json = json.dumps([i.tolist() for i in embedding])
    insert_sql = """
                insert into t_user_info (user_code, group_id, username,user_remark, embedding, create_time)
                values ({},'{}','{}','{}','{}', now());
                """.format(user_code, group_id, user_name, '', embedding_json)
    count = execute_db(insert_sql)

def update_user_info(user_code, embedding):
    embedding_json = json.dumps([i.tolist() for i in embedding])
    update_sql = """
                update t_user_info set embedding='{}' where user_code='{}';
                """.format(embedding_json, user_code)
    count = execute_db(update_sql)


def get_group_info(group_name):
    sql_s = "select group_id, group_name, group_remark from t_group where group_name='{group_name}';".format(
        group_name=group_name)
    group_info = query_db(sql_s)
    if group_info and len(group_info) > 0:
        return {
            'groupId': group_info[0][0],
            'groupName': group_info[0][1],
            'groupRmark': group_info[0][2]
        }
    else:
        return None

def get_group_info_by_id(group_id):
    sql_s = "select group_id, group_name, group_remark from t_group where group_id='{group_id}';".format(
        group_id=group_id)
    group_info = query_db(sql_s)
    if group_info and len(group_info) > 0:
        return {
            'groupId': group_info[0][0],
            'groupName': group_info[0][1],
            'groupRmark': group_info[0][2]
        }
    else:
        return None


def insert_group_info(group_name, group_remark):
    insert_sql = """insert into t_group (group_name, group_remark)values ('{}','{}');""".format(group_name, group_remark)
    count = execute_db(insert_sql)

def delete_group_info(group_id):
    delete_sql = """delete from t_group where group_id='{}'""".format(group_id)
    count = execute_db(delete_sql)


def update_group_info(group_name, group_remark):

    update_sql = """
                update t_group set group_remark='{}' where group_name='{}';
                """.format(group_remark, group_name)
    count = execute_db(update_sql)




def post_server(post_url, data):
    res = requests.post(url=post_url, data=data)
    res_val = json.loads(res.text)
    if res_val.get('code') == 1:
        return None, res_val.get('data')
    else:
        return (90099, '系统异常'), None

def img_resize(img):
    h, w, c = img.shape
    if h > 300 or w > 300:
        h_scale = 300
        w_scale = int(h_scale / h * w)
        img = cv2.resize(img, (w_scale, h_scale))
    return img

def create_success(message, data = None):
    if data:
        pass
    return {
                'code' : 0,
                'msg' : message
            }

def create_error(error, message):
    return {
                'code' : -1,
                'msg' : message,
                "errorCode": error[0],
                "errorMsg": error[1]
            }

# 人脸注册api
# 传入三张图片，后台对人脸进行特征提取，存储用户id和特征，返回是否注册成功
@app.post("/face_register")
async def faceRegister(userCode: str = Form(...), userName: str = Form(...), groupId: str = Form(...), file: List[str] = Form(...)):
    try:
        # userCode = request.form.get("userCode")
        # userName = request.form.get("userName")
        # groupId = request.form.get("groupId")
        # file = request.form.getlist("file")

        is_validate_pass = True
        error_msg = []

        if file is None or len(file) != 3:
            is_validate_pass = False
            error_msg.append('图片不等于三张')
            return create_error((90007, '图片不等于三张'), '注册信息错误')

        file_count = 0
        for f in file:
            if f:
                file_count += 1

        if file_count != 3:
            return create_error((90007, '图片不等于三张'), '注册信息错误')

        if groupId is None:
            is_validate_pass = False
            error_msg.append('缺少groupid')
            return create_error((90001, '缺少groupid'), '注册信息错误')

        if userName is None:
            is_validate_pass = False
            error_msg.append('缺少userName')
            return create_error((90002, '缺少userName'), '注册信息错误')

        if userCode is None:
            is_validate_pass = False
            error_msg.append('缺少userCode')
            return create_error((90003, '缺少userCode'), '注册信息错误')


        if is_validate_pass:

            userInfo = get_user_info(userCode)
            if userInfo is not None:
                return create_error((90010, '该用户已注册'), '注册信息错误')
            i = 1
            # 将base64编码图片变成cv2可读图片

            img_dirs = []
            for item in file:

                base64Image = item.replace("%2B", "+").replace("%3D", "=").replace("%2F", "/")
                img = base64.b64decode(base64Image)
                img_array = np.frombuffer(img, np.uint8)
                ori_img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
                ori_img = img_resize(ori_img)

                dir = './imgs/' + userCode
                if not os.path.exists(dir):
                    os.makedirs(dir)
                num = datetime.now().strftime("%m-%d-%Y-%H-%M-%S") + '-' + str(i)
                # 写原图
                cv2.imwrite(dir + '/' + num + ".jpg", ori_img)
                error, ali_img, box_val = retina_model.detect_face(ori_img)
                # ali_img = np.transpose(ali_img, (1, 2, 0))
                # ali_img = cv2.cvtColor(ali_img, cv2.COLOR_BGR2RGB)
                # print(ali_img.shape)
                if error:
                    return create_error(error, '注册失败')
                # print(ori_img)
                # print(type(ori_img))
                # dir = './imgs/' + userCode + userName
                # num = datetime.now().strftime("%m-%d-%Y-%H-%M-%S") + '-' + str(i)
                # if not os.path.exists(dir):
                #     os.makedirs(dir)
                # # w, h = ori_img.shape[:2]
                # # if w < 112 or h < 112:
                # #     is_validate_pass = False
                # #     error_msg.append('图像不清晰')
                # #     break
                #
                # cv2.imwrite(dir + '/' + num + ".jpg", ori_img)
                img_dirs.append(ali_img)
                i += 1
        # print(file)
        # print(len(file))
        if is_validate_pass:
            try:
                face_embeddings = []
                for (i, imagePath) in enumerate(img_dirs):
                    image = imagePath
                    # convert face to RGB color
                    nimg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    nimg = cv2.resize(nimg, (112, 112))
                    nimg = np.transpose(nimg, (2, 0, 1))
                    # Get the face embedding vector
                    error, face_embedding = face_feature.get_feature(nimg)
                    if error:
                        return create_error(error, '注册失败')
                    face_embeddings.append(face_embedding)

                error, embeddings = register_pkl_embedding(userCode, userName, groupId, face_embeddings)
                # print(embeddings)
                # print(embeddings.shape)
                if error:
                    return create_error(error, '注册失败')
                else:
                    # 插入新注册用户
                    url_str = faiss_server_url + 'face_add'
                    data = {'groupId': groupId, 'userCode': userCode, 'embedding': json.dumps([i.tolist() for i in embeddings])}

                    error, result_data = post_server(url_str, data)
                    if error:
                        return create_error((90099, '系统异常'), '注册失败')
                    # 插入用户信息
                    insert_user_info(userCode, groupId, userName, embeddings)
                    # add_faiss_embedding(groupId, userCode, embeddings)
                    ###### 更新faiss
                    ###### refresh_faiss(groupId, get_db())
                    return create_success('注册成功')

            except Exception as e:
                traceback.print_exc()           ###
                return create_error((90099, '系统异常'), '注册失败')
        else:
            return create_error((90099, '系统异常'), '注册失败')
    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '注册失败')

#人脸认证1:1
#传入一张图片和用户id，与该用户注册的人脸进行对比，返回是否匹配、外接矩形
@app.post("/face_auth_one")
async def faceAuthOne(userCode: str = Form(...), file: str = Form(...)):
    print(time.time())
    try:
        if file is None or file == '':
            return create_error((90008, '缺少图片'), '人脸认证错误')

        if userCode is None:
            return create_error((90003, '缺少userCode'), '人脸认证错误')
        t1 = time.time()
        userInfo = get_user_info(userCode)
        if userInfo is None:
            return create_error((90009, '该用户未注册'), '人脸认证错误')

        base64Image = file.replace("%2B", "+").replace("%3D", "=").replace("%2F", "/")
        img = base64.b64decode(base64Image)
        img_array = np.frombuffer(img, np.uint8)
        ori_img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
        ori_img = img_resize(ori_img)
        error, ali_img, box_val = retina_model.detect_face(ori_img)
        if error:
            return create_error(error, '人脸认证错误')
        try:
            frame = cv2.cvtColor(ali_img, cv2.COLOR_BGR2RGB)
            frame = cv2.resize(frame, (112, 112))
            frame = np.transpose(frame, (2, 0, 1))
            error, face_embedding = face_feature.get_feature(frame)
            if error:
                return create_error(error, '人脸认证错误')
            error, result, dis, ou_dis = one2onne_embedding(userCode, userInfo['embedding'], face_embedding)
            print(ou_dis)
            if error:
                return create_error(error, '人脸认证错误')
            return {
                "code": 0,
                "data": {
                    "match": result,
                    "box": box_val,
                    "cosdis": dis,
                    "eluerdis":[i.tolist() for i in ou_dis]
                }
            }

        except Exception as e:
            traceback.print_exc()  ###
            return create_error((90099, '系统异常'), '注册失败')
    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '注册失败')

#人脸是否注册API
#传入一张图片，与人脸注册库中的人脸进行对比，返回是否注册、外接矩形
@app.post("/face_is_register")
async def faceIsRegister(userCode: str = Form(...)):
    try:
        if userCode is None:
            return create_error((90003, '缺少userCode'), '人脸查询错误')
        userInfo = get_user_info(userCode)
        result = False
        if userInfo is not None:
            result = True

        return {
            "code" : 0,
            "data" : {
                "register" : result
            }
        }
    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '人脸查询错误')


#人脸更新API
#传入三张图片，后台对人脸进行特征提取，替换用户原来的特征，返回是否更新成功
@app.post("/face_update")
async def faceUpdate(userCode: str = Form(...), userName: str = Form(...), groupId: str = Form(...), file: List[str] = Form(...)):
    try:
        is_validate_pass = True
        error_msg = []

        if file is None or len(file) != 3:
            is_validate_pass = False
            error_msg.append('图片不等于三张')
            return create_error((90007, '图片不等于三张'), '注册信息错误')

        file_count = 0
        for f in file:
            if f:
                file_count += 1

        if file_count != 3:
            return create_error((90007, '图片不等于三张'), '注册信息错误')

        if groupId is None:
            is_validate_pass = False
            error_msg.append('缺少groupid')
            return create_error((90001, '缺少groupid'), '注册信息错误')

        if userName is None:
            is_validate_pass = False
            error_msg.append('缺少userName')
            return create_error((90002, '缺少userName'), '注册信息错误')

        if userCode is None:
            is_validate_pass = False
            error_msg.append('缺少userCode')
            return create_error((90003, '缺少userCode'), '注册信息错误')

        if is_validate_pass:

            userInfo = get_user_info(userCode)
            if userInfo is None:
                return create_error((90009, '该用户未注册'), '注册信息错误')
            i = 1
            # 将base64编码图片变成cv2可读图片
            img_dirs = []
            for item in file:

                base64Image = item.replace("%2B", "+").replace("%3D", "=").replace("%2F", "/")
                img = base64.b64decode(base64Image)
                img_array = np.frombuffer(img, np.uint8)

                ori_img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
                ori_img = img_resize(ori_img)
                # print(ori_img.shape)

                dir = './imgs/' + userCode
                if not os.path.exists(dir):
                    os.makedirs(dir)
                num = datetime.now().strftime("%m-%d-%Y-%H-%M-%S") + '-' + str(i)
                # 写原图
                cv2.imwrite(dir + '/' + num + ".jpg", ori_img)
                # print('----------------------------')
                # print(img_array)
                error, ali_img, box_val = retina_model.detect_face(ori_img)
                # print(ali_img)
                # ali_img = np.transpose(ali_img, (1, 2, 0))
                # ali_img = cv2.cvtColor(ali_img, cv2.COLOR_BGR2RGB)
                if error:
                    return create_error(error, '更新失败')
                # print(ori_img)
                # print(type(ori_img))


                # w, h = ori_img.shape[:2]
                # if w < 112 or h < 112:
                #     is_validate_pass = False
                #     error_msg.append('图像不清晰')
                #     break

                # cv2.imwrite(dir + '/' + num + ".jpg", ali_img)
                # img_dirs.append(dir + '/' + num + ".jpg")
                img_dirs.append(ali_img)
                i += 1

        # print(file)
        # print(len(file))
        if is_validate_pass:
            try:
                face_embeddings = []
                for (i, image) in enumerate(img_dirs):
                    nimg = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    nimg = cv2.resize(nimg, (112, 112))
                    nimg = np.transpose(nimg, (2, 0, 1))
                    error, face_embedding = face_feature.get_feature(nimg)
                    if error:
                        return create_error(error, '更新失败')
                    face_embeddings.append(face_embedding)

                error, face_embedding_lst = re_register_pkl_embedding(userCode, userName, groupId, face_embeddings)
                if error:
                    return create_error(error, '更新失败')
                else:
                    update_user_info(userCode, face_embedding_lst)
                    # 更新faiss
                    url_str = faiss_server_url + 'face_reload'
                    data = {'groupId': groupId}
                    error, result_data = post_server(url_str, data)
                    if error:
                        return create_error(error, '更新失败')
                    #refresh_faiss(groupId, get_db())
                    return create_success('更新成功')

            except Exception as e:
                traceback.print_exc()  ###
                return create_error((90099, '系统异常'), '更新失败')
        else:
            # error_msg
            return create_error((90099, '系统异常'), '更新失败')

    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '注册失败')


#人脸认证API（1:N）
#传入一张图片，与人脸注册库中的人脸进行对比，返回用户id、外接矩形
@app.post("/face_auth_n")
async def faceAuthN(groupId: str = Form(...), file: str = Form(...)):
    try:

        if file is None  or file == '':
            return create_error((90008, '缺少图片'), '人脸认证错误')

        if groupId is None:
            return create_error((90001, '缺少groupid'), '人脸认证错误')

        base64Image = file.replace("%2B", "+").replace("%3D", "=").replace("%2F", "/")
        img = base64.b64decode(base64Image)
        img_array = np.frombuffer(img, np.uint8)
        ori_img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)
        ori_img = img_resize(ori_img)
        error, ali_img, box_val = retina_model.detect_face(ori_img)
        if error:
            return create_error(error, '人脸认证错误')
        try:
            frame = cv2.cvtColor(ali_img, cv2.COLOR_BGR2RGB)
            frame = cv2.resize(frame, (112, 112))
            frame = np.transpose(frame, (2, 0, 1))
            # Get the face embedding vector

            error, face_embedding = face_feature.get_feature(frame)
            if error:
                return create_error(error, '人脸认证错误')

            url_str = faiss_server_url + 'face_search'
            data = {'groupId': groupId, 'embedding': json.dumps(face_embedding.tolist())}

            error, result_data = post_server(url_str, data)
            # error, userCode,I,D = faiss_search(groupId, ali_img, embedding_model)
            if error:
                return create_error(error, '人脸认证错误')
            userCode = result_data.get('userCode')
            I = result_data.get('I')
            D = result_data.get('D')
            user_info = get_user_info(userCode)
            return {
                "code": 0,
                "data": {
                    "userCode": userCode,
                    "box": box_val,
                    "index":I,
                    "prob":D
                }
            }
        except Exception as e:
            traceback.print_exc()  ###
            return create_error((90099, '系统异常'), '人脸认证错误')
    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '人脸认证错误')


@app.get("/")
def index():
    print(time.time())
    time.sleep(10)
    return {"message": "oaclockin service is runing"}


# 注册group_id
# 传入group_name,remark，返回是否注册成功
@app.post("/facedb_register")
async def groupRegister(groupname: str = Form(...), groupremark: str = ''):
    try:
        groupName = groupname
        groupRmark = groupremark

        is_validate_pass = True

        if groupName is None:
            return create_error((90001, '缺少groupName'), '注册信息错误')

        groupInfo = get_group_info(groupName)
        if groupInfo is not None:
            return create_error((90010, '底库已注册'), '注册信息错误')

        if is_validate_pass:
            insert_group_info(groupName, groupRmark)
            return create_success('注册成功')

    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '注册失败')

@app.post("/facedb_delete")
async def groupDelete(groupId: str = Form(...)):
    try:
        delete_group_info(groupId)
        return create_success('删除成功')
    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '删除失败')

# groupId更新
@app.post("/group_update")
async def groupUpdate(groupname: str = Form(...), groupremark: str = None):
    try:
        groupName = groupname
        groupRmark = groupremark

        is_validate_pass = True

        if groupName is None:
            return create_error((90001, '缺少groupName'), '底库更新错误')

        if is_validate_pass:

            groupInfo = get_group_info(groupName)
            if groupInfo is None:
                return create_error((90009, 'group未注册'), '底库更新错误')
            update_group_info(groupName, groupRmark)
            return create_success('group更新成功')

    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), '注册失败')

#groupId是否注册
@app.post("/facedb_is_register")
async def groupIsRegister(groupId: str = Form(...)):
    try:
        if groupId is None:
            return create_error((90001, '缺少groupid'), '底库查询错误')

        groupInfo = get_group_info_by_id(groupId)
        result = False
        if groupInfo is not None:
            result = True

        #### 此处调用方法

        return {
            "code" : 0,
            "data" : {
                "register" : result
            }
        }
    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), 'group查询错误')

#输入groupName，返回整条数据
@app.post("/facedb_search")
async def groupSearch(groupId: str = Form(...)):
    try:
        if groupId is None:
            return create_error((90003, '缺少groupId'), '底库查询错误')

        groupInfo = get_group_info_by_id(groupId)

        if groupInfo is None:
            return create_error((90012, '底库未注册'), '底库查询错误')
        else:
            return {
                "code": 0,
                "data": groupInfo
            }

    except Exception as e:
        traceback.print_exc()  ###
        return create_error((90099, '系统异常'), 'group查询错误')



# export LC_ALL=en_US.utf-8
# export LANG=en_US.utf-8

# uvicorn api_oa_face_fastapi:app --reload --host 0.0.0.0 --port 17104

# uvicorn api_oa_face_fastapi:app --host 0.0.0.0 --port 17104 --workers 4